/*
    Author: Christopher Diggins
    License: MIT Licence 1.0
    Website: http://www.cdiggins.com
*/

#ifndef CVML_COMPILER_HPP
#define CVML_COMPILER_HPP

struct Compiler 
{
    ByteCode* bytecode;	
	lls::Module* module;
	heron_ast::HeronAst* ast;
	LuaMachine luavm;
	int oldsize;
	int originalsize;
	int pass;

	Compiler() 
	{
		bytecode = NULL;
		module = NULL;
		ast = NULL;
		oldsize = 0;
		originalsize = 0;
		pass = 0; 
	}

	bool Compile()
	{
		try
		{
			if (!LoadFile()) {
				return false;
			}
			HeronToLLScheme0();
			LLScheme0ToLua();
			RenameVars();
			LLScheme0ToLLScheme1();
			RemoveNames();
			LLScheme1ToCVML();
			PrintFinal();
			return true;
		}
		catch (CVMLException e) 
		{
			Log(String("Error occured: ") + e.msg);
			return false;
		}
		catch (...)
		{
			Log("Unidentified error occured");
			return false;
		}	
	}
	
	bool LoadFile()
	{
		ast = new heron_ast::HeronAst();		
		return ast->LoadFile(sInFile.c_str());
	}

	~Compiler() {
		if (module) delete module;
		if (bytecode) delete bytecode;
		if (ast) delete ast;
	}
 
	String GetTmpFileName() {
		return sOutFolder + "\\" + sOutFileName + ToStr(pass) + ".tmp";
	}

	double ComputeEntropy(Binary& binary) {
		std::vector<int> cnts(256);
		int total = binary.code.Count();
		if (total == 0) 
			return 0.0;
		for (int i=0; i < total; ++i) {
			int n = binary.code[i];
			assert(n >= 0 && n < 256);
			cnts[n] += 1;
		}
		double entropy = 0.0;
		for (int i=0; i < 256; ++i) {
			double prob = (double)cnts[i] / (double)total;
			if (prob > 0.0) {
				double log2prob = log(prob) / log(2.0);
				entropy -= prob * log2prob;
			}
		}
		return entropy;
	}

	void SaveByteCodeAsBinary(ConstString bin_filename) {
		Binary bin;
		bytecode->ConvertToBinary(bin);
		FileWritingStream stream(bin_filename);
		bin.Serialize(stream);
	}

	void PrintLLScheme(ConstString sMsg) {
		if (bSummaryOnly)
			return;

		Log("Pass " + ToStr(pass) + ": " + sMsg);		
		if (bOutputIntermediateCode) {
			String s = lls::PrettyPrint(module);
			Log(s);
		}
		++pass;
	}

	void Print(FILE* f, Instruction* i, int& offset)
	{
		if (i->operand >= 0)
			fprintf(f, "  %d: %s(%x);\t %s\n", offset, OpCodeToPChar(i->opcode), i->operand, i->comment.c_str());
		else
			fprintf(f, "  %d: %s;\t %s\n", offset, OpCodeToPChar(i->opcode), i->comment.c_str());
		offset += GetOpCodeWidth(i->opcode);
	}

	void Print(FILE* f, Subroutine* sub, int& offset)
	{
		fprintf(f, "%s // %x\n", sub->name.c_str(), sub);
		fprintf(f, "{\n");
		for (Instruction* tmp = sub->first; tmp != NULL; tmp = tmp->next) 
			Print(f, tmp, offset);
		fprintf(f, "}\n");
	}

	void Print(FILE* f, ByteCode* bytes)
	{
		int offset = 0;
		fprintf(f, "CONSTANTS\n{\n");
		for (int i=0; i < bytes->GetNumConstants(); ++i) {
			fprintf(f, "  %d = \"%s\"\n", i, bytes->GetConstant(i).c_str());
		}
		fprintf(f, "}\nTotal constant pool size %d\n\n", bytes->GetTotalConstantSize());
		fprintf(f, "CODE\n");
		for (int i=0; i < bytes->GetNumSubs(); ++i) {
			fprintf(f, "sub(%d): ", i);
				Print(f, bytes->GetSub(i), offset);
		}
	}

	void Log(ConstString s) 
	{
		printf("%s\n", s.c_str());
	}

	void OutputStats() {
		if (bNoStats) 
			return;

		Log("Bytecode size = " + ToStr(bytecode->GetTotalSize()));
		Binary binary;
		bytecode->ConvertToBinary(binary);		
		
		if (!bShort) {
			Log("Binary size = " + ToStr(binary.GetTotalSize()));
			Log("Entropy = " + ToStr(ComputeEntropy(binary)));
		}

		/*
		Log("Stats:");
		int newsize = bytecode->GetTotalSize();
		if (bFirst) 
			originalsize = newsize;
		Log("Size = " + ToStr(newsize));
		int compression = oldsize - newsize;
		oldsize = newsize;
		double ratio = (double)compression / (double)oldsize;
		if (!bFirst && !bFinal) {
			Log("Incremental compression amount = " + ToStr(compression));
			Log("Incremental compression percent = " + ToStr(ratio * 100.0) + "%");		
		}
		int total_compression = originalsize - newsize;
		double total_ratio = (double)total_compression / (double)originalsize;
		if (!bFirst) {
			Log("Total compression amount = " + ToStr(total_compression));
			Log("Total compression percent = " + ToStr(total_ratio * 100.0) + "%");		
		}
		oldsize = newsize;
		Binary binary;
		bytecode->ConvertToBinary(binary);		
		Log("Binary size = " + ToStr(binary.GetTotalSize()));
		Log("Entropy = " + ToStr(ComputeEntropy(binary)));
		*/
	}

	void PrintByteCode(ConstString sMsg) {
		if (bSummaryOnly)
			return;

		Log("Pass " + ToStr(pass) + ": " + sMsg);

		if (bOutputIntermediateCode) {
			Print(stdout, bytecode);
		}
		if (bOutputIntermediateStats) {
			OutputStats();
		}

		++pass;
	}
	
	void PrintBinaryCode(FILE* f, Binary* bin)
	{
		u1* begin = bin->code.Begin();
		u1* end = bin->code.End();
		int curfxn = 0;
		int offset(0);
		
		while (begin != end) 
		{         
			int n = begin - bin->code.Begin();
			if (n >= bin->fxn_indexes[curfxn]) {
				printf("\nSubroutine %d\n", curfxn++);
			}

			// compute opcode width
			OpCodeEnum opcode = (OpCodeEnum)(*begin++);
			u4 operand = 0;            
			switch (GetOpCodeWidth(opcode)) {
				case 0:
				case 1:
					operand = 0;
					break;
				case 2:
					operand |= *begin++;
					break;
				case 3:
					operand |= *begin++ << 8;
					operand |= *begin++;
					break;
				case 5:
					operand |= *begin++ << 24;
					operand |= *begin++ << 16;
					operand |= *begin++ << 8;
					operand |= *begin++;
					break;
				default:
					Error("unexpected opcode width");
			}
			Assert(begin <= end);
			Instruction i(opcode, operand, "");
			Print(f, &i, offset);
		}            
	}

	void HeronToLLScheme0() {
		module = heron_to_lls::Convert(ast->root);
		PrintLLScheme("Initial s-expression representation");
		lls::PlaceReturnsInTailPosition(module);
		PrintLLScheme("Return statements converted to tail position");		
		if (bSaveLLSFile) {
			String s = lls::PrettyPrint(module);
			PrintStringToFile(s, sOutFile + ".lls");
		}
	}

	void RenameVars() {
		lls::UniquelyNameVars(module);
		PrintLLScheme("Uniquely re-named variables");
	}

	void LLSToTable() {
		luavm += "ast = {\n";
		lls::LLSchemeToLuaTable tableMaker(luavm);
		tableMaker(module);
		luavm += "}\n";
	}

	void LLSToLua(bool bPreprocess) {
		lls::LLSchemeToLua luaMaker(luavm, bPreprocess);
		luaMaker.ModuleToLua(module);
	}

	lls::Module* LLSFromLua() {
		lls::LLchemeFromLua sexprMaker(luavm);
		lls::SExpr* ast = sexprMaker.FromLuaMachine();
		lls::Module* mod = ast->GetChild(0)->As<lls::Module>();
		Assert(mod != NULL);
		return mod;
	}

	void PrintLua(ConstString sFile) {
		if (bSaveLuaFile) {
			String s = luavm.GetPrettyText();
			PrintStringToFile(s, sFile);
		}
	}

	void PerformPreprocess() {
		if (!bRunPreprocess)
			return;
		if (!luavm.GlobalExists("preprocess")) 
			return;
		if (!bSummaryOnly) 
			Log("Performing preprocess step");
		luavm.Eval("preprocess()");
		lls::Module* tmp = LLSFromLua();
		delete module;
		module = tmp;
		if (bSaveLLSFile) {
			String s = lls::PrettyPrint(module);
			PrintStringToFile(s, sOutFile + ".2.lls");
		}
		luavm.Reset();
		LLSToLua(false);
		PrintLua(sOutFile + ".2.lua");
	}

	void LLScheme0ToLua() {
		if (!bSummaryOnly)
			Log("Generating Lua table from AST");
		LLSToTable();
		if (!bSummaryOnly) 
			Log("Generating Lua program from AST");
		LLSToLua(true);
		PrintLua(sOutFile + ".lua");
		PerformPreprocess();
	}

	void Run() {
		luavm.Eval("main()");
	}

	void LLScheme0ToLLScheme1() {
		lls::RemoveImplicitScopes(module);
		PrintLLScheme("Removed implicit scopes");
		//lls::MarkFinals(module);
		//PrintLLScheme("Marked finals");
	}

	void RemoveNames() {
		lls::SExpr* old_module = module;
		module = lls::RemoveNames(module, bVerboseNameRemoval);
		delete old_module;
		PrintLLScheme("Removed names");
	}

	void LLScheme1ToCVML() {
		bytecode = new ByteCode();
		lls::ModuleToCVML(module, bytecode);

		if (!bSummaryOnly) {
			Log("Raw Byte-code");
			if (bOutputIntermediateCode) {
				Print(stdout, bytecode);
			}
			OutputStats();
		}

		if (bGenerateParameterizedSubs) {
			GenerateParameterizedSubs(bytecode);
			PrintByteCode("Generated parameterized subroutines");
		}

		if (bGenerateSubs) {
			GenerateSubs(bytecode);
			PrintByteCode("Generated new subroutines");
		}

		if (bInline) {
			InlineSubs(bytecode);
			PrintByteCode("Inlined single-use and small subroutines");
		}
	
		if (bMakeTailCalls) {
			MakeTailCalls(bytecode);
		}

		ResolveSubroutines(bytecode);
		//PrintByteCode("Resolved subroutines");

		if (bCompaction) {
			ReplaceDuplicateSubs(bytecode);
			Peephole(bytecode);
			CompactInstructions(bytecode);
			PrintByteCode("Compacted instructions");
		}		
	}

	void PrintFinal() {
		OutputStats();
		FILE* f = fopen(sOutFile.c_str(), "w");	
		Print(f, bytecode);
		fclose(f);
	}
};

#endif
 